﻿using System;
using System.Collections.Generic;
using System.Linq;
using Bettzueche.RLLibrary.Types;
using Bettzueche.RLLibrary.Functions;
using Bettzueche.RLLibrary.Util;

namespace Bettzueche.RLLibrary.Component {
    /// <summary>
    /// Generische Implementierung eines Sarsa(λ) Agenten mit ersetzenden Traces.<para>
    /// Es müssen minimal die Eigenschaften Environment, QFunktion, Policy gesetzt werden.</para><para>
    /// Wenn nicht anders gesetzt, wird <see cref="SarsaLambda{S,A}.DefaultEFunc"/> als EFunction genutzt.</para>
    /// </summary>
    /// <remarks>
    /// Diese Implementierung nutzt TraceSet, um (s,a) Tupel einer Episode zu speichern und so Updates effizient umzusetzen.<para>
    /// Wesentlicher Update-Schritt in DoStep: Q(s,a) += α·δ·e(s,a)</para>
    /// </remarks>
    public class SarsaLambda<S,A> : GenericAgent<S,A>
        where S : IState
        where A : IAction
    {
        /// <summary>
        /// Gibt die insgesamt erhaltene Vergeltung an.<para>
        /// Wird mit Reset() auf 0.0 gesetzt.</para>
        /// </summary>
        /// <value>
        /// Die neue Gesamtentlohnung.
        /// </value>
        /// <remarks>Es werden 'rewards' von Environment.DoAction(), die in DoStep() aufgerufen wird, aufsummiert.</remarks>
        public double RewardTotal { get; protected set; }



        #region Konstruktor(en)

        /// <summary>
        /// Initializes a new instance of the <see cref="SarsaLambda{S, A}"/> class.<para>
        /// Damit der Algorithmus funktioniert , müssen Environment, QFunktion und Policy gesetzt werden!
        /// Für EFunction wird eine Standardimplementierung verwendet, falls nicht anders gesetzt wird.</para>
        /// </summary>
        internal SarsaLambda() {
            // Path = new TraceSet<SATupel<S, A>>();
            // Default Impls setzen, wo's nötig ist:
            EFunction = new DefaultEFunc();
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="SarsaLambda{S, A}"/> class.<para>
        /// Für EFunction wird eine Standardimplementierung verwendet, falls nicht anders gesetzt wird.</para>
        /// </summary>
        /// <param name="env">The env.</param>
        /// <param name="qFunc">The q function.</param>
        /// <param name="policy">The policy.</param>
        public SarsaLambda(IEnvironment<S,A,double> env, IStateValueFunction<S,A> qFunc, GenericExplorationPolicy<A> policy) 
            : this()
        {
            this.Environment = env;
            this.QFunction = qFunc;
            this.Policy = policy;
        }

        #endregion


        #region IAgent Members

        #region Properties

        /// <summary>
        /// Gets or sets the exploration policy.
        /// <para>(Optional)</para>
        /// </summary>
        /// <value>
        /// The exploration policy.
        /// </value>
        public override GenericExplorationPolicy<A> Policy {
            get;
            set;
        }

        /// <summary>
        /// Not Supported.
        /// </summary>
        /// <value>
        /// Not Supported.
        /// </value>
        public override StochasticPolicy Pi {
            get {
                throw new NotImplementedException();
            }
            set {
                throw new NotImplementedException();
            }
        }

        /// <summary>
        /// Gets or sets the environment.
        /// </summary>
        /// <value>
        /// The environment.
        /// </value>
        public override IEnvironment<S, A, double> Environment {
            get;
            set;
        }

        double _alpha;
        /// <summary>
        /// Gets or sets the learn rate ('alpha').
        /// <para>(Optional)</para>
        /// </summary>
        /// <value>
        /// The learn rate from 0 to 1. Truncated, if not.
        /// </value>
        public override double LearnRate {
            get { return _alpha; }
            set { _alpha = Math.Max(0, Math.Min(1.0, value)); }
        }

        double _gamma;
        /// <summary>
        /// Gets or sets the discount factor ('gamma') for this Agent and his <c>Efunction</c>.
        /// </summary>
        /// <value>
        /// The discount factor from 0 to 1. Truncated, if not.
        /// </value>
        public override double DiscountFactor {
            get { return _gamma; }
            set { 
                _gamma = Math.Max(0, Math.Min(1.0, value));
                EFunction.DiscountFactor = _gamma;
            }
        }

        double _lambda;
        /// <summary>
        /// Gets or sets the trace decay ('lambda') for this Agent and his <c>Efunction</c>.
        /// <para>(Optional)</para>
        /// </summary>
        /// <value>
        /// The trace decay from 0 to 1. Truncated, if not.
        /// </value>
        public override double TraceDecay {
            get { return _lambda; }
            set { 
                _lambda = Math.Max(0, Math.Min(1.0, value));
                EFunction.TraceDecay = _lambda;
            }
        }

        /// <summary>
        /// Not Supported
        /// </summary>
        /// <value>
        /// Not Supported
        /// </value>
        /// <exception cref="NotImplementedException" />
        public override IValueFunction<S> VFunction {
            get {
                throw new NotImplementedException();
            }
            set {
                throw new NotImplementedException();
            }
        }

        /// <summary>
        /// Gets or sets the Q function.
        /// </summary>
        /// <value>
        /// The Q function.
        /// </value>
        public override IStateValueFunction<S, A> QFunction {
            get;
            set;
        }

        /// <summary>
        /// Gets or sets the Eligibility function.
        /// </summary>
        /// <value>
        /// The E function.
        /// </value>
        public override IEligibilityFunction<A, S> EFunction {
            get;
            set;
        }

            #endregion

        /// <summary>
        /// Überprüft, ob notwendige Eigenschaften gesetzt sind.
        /// </summary>
        /// <param name="initialState">wird ignoriert. Der Zustand wird stattdessen intern von Environment abgefragt.</param>
        /// <param name="specifications">wird ignoriert</param>
        /// <exception cref="RLLibraryException">Wenn Environment, QFunktion, Policy oder EFunction null sind.</exception>
        public override void Init(S initialState, string specifications) {
            //Damit der Algorithmus funktioniert , müssen Environment, QFunktion und Policy gesetzt werden!
            //Für EFunction wird eine Standardimplementierung verwendet, falls nicht anders gesetzt wird.
            if (Environment == null)
                throw new RLLibraryException("Environment darf nicht null sein!");
            if (QFunction == null)
                throw new RLLibraryException("QFunction darf nicht null sein!");
            if (Policy == null)
                throw new RLLibraryException("Policy darf nicht null sein!");
            if (EFunction == null)
                throw new RLLibraryException("EFunction darf nicht null sein!");
        }

        /// <summary>
        /// Führt die angegebene Aktion aus und aktualisiert daraufhin <c>QFunction</c>.<para>
        /// Umsetzung eines Iterationsschritts im Sarsa(λ) Algorithmus.</para>
        /// </summary>
        /// <param name="action">Die auszuführende Aktion.</param>
        /// <returns>nächste gemäß aktueller Strategie auszuführende Aktion.</returns>
        /// <remarks>
        /// Bedingungen: Die Efunction muss Update implementieren. Die QFunction muss Indexer Q[s,a] und Add(Efunction, double)
        /// implementieren.
        /// </remarks>
        public override A DoStep(A action) {
            S state = Environment.State;
            S nextState;
            double reward = Environment.DoAction(action, out nextState);
            RewardTotal += reward;

            A nextAction = Policy.GetAction(QFunction.GetEstimates(nextState));
            double delta = reward + _gamma * QFunction[nextState, nextAction] - QFunction[state, action];
            EFunction.Update(action, state); // E = γ·λ·E; e_sx = 0; e_sa = 1;
            QFunction.Add(EFunction, LearnRate * delta); // W = W + α·δ·E
            /*
            QFunction.Values = QFunction.Values + (LearnRate * delta) * EFunction.Eligibilities;
            Path.Add(new SATupel<S, A>(state, action)); // add to visited (s,a)'s
            foreach (var tupel in Path) {
                // do update
            }
            */
            return nextAction;
        }

        ///// <summary>
        ///// Führt die angegebene Aktion aus und aktualisiert daraufhin <c>QFunction</c>.<para>
        ///// Umsetzung eines Iterationsschritts im Sarsa(λ) Algorithmus. Benutzt ersetzende Traces.</para>
        ///// </summary>
        ///// <param name="action">Die auszuführende Aktion.</param>
        ///// <returns>nächste gemäß aktueller Strategie auszuführende Aktion.</returns>
        ///// <remarks>Die Methode führt ein Update der QFunction und EFunction Werte für alle 
        ///// (s,a)-Tupel aus <c>Path</c> aus.</remarks>
        //public override A DoStep(A action) {
        //    S state = Environment.State;
        //    S nextState;
        //    double reward = Environment.DoAction(action, out nextState);
        //    Path.Add(new SATupel<S, A>(state, action)); // add to visited (s,a)'s
        //    RewardTotal += reward;

        //    A nextAction = Policy.GetAction(QFunction.GetEstimates(nextState));
        //    double delta = reward + this.DiscountFactor * QFunction[nextState, nextAction] - QFunction[state, action]; // δ

        //    double decayRate = this.DiscountFactor * this.TraceDecay; // γ·λ
        //    double qUpdateFactor = this.LearnRate * delta; // α·δ

        //    var removeFromPath = new HashSet<SATupel<S, A>>();
        //    foreach (var tupel in Path.Trace) {
        //        S s = tupel.Item1;
        //        A a = tupel.Item2;
        //        double e_i;
        //        if (s.Equals(state))
        //            e_i = (a.Equals(action)) ? 1.0 : 0.0;
        //        else {
        //            e_i = EFunction[a, s] * decayRate;
        //            if (e_i < this.MinE) { // üblicherweise ein oder kein Element aus Path
        //                e_i = 0.0;
        //                removeFromPath.Add(tupel);
        //            }
        //        }
        //        if (e_i != 0.0) {
        //            EFunction[a, s] = e_i;
        //            QFunction[s, a] += qUpdateFactor * EFunction[a, s];
        //        }
        //    }
        //    foreach (var tupel in removeFromPath) {
        //        Path.Remove(tupel);
        //    }

        //    return nextAction;
        //}

        /// <summary>
        /// Runs the episode.
        /// </summary>
        /// <param name="maxActionCount">The maximum action count.</param>
        /// <returns>Anzahl ausgeführter Schritte</returns>
        public override int RunEpisode(int maxActionCount) {
            // alles Initialisiert?
            A action = Policy.GetAction(QFunction.GetEstimates(Environment.State));
            int stepCount = 0;
            while ((!Environment.IsFinal) && stepCount < maxActionCount) {
                action = this.DoStep(action);
                stepCount++;
            }
            // Falls einer der Endzustände erreicht wurde, wird die Lernepisode automatisch beendet.
            //if (finals.Contains<int>(environment.State))
            //    resetEpsiode();
            return stepCount;
        }

        public override int RunEpisode(params S[] finalStates) {
            return RunEpisode(Int32.MaxValue, finalStates);
        }

        /// <summary>
        /// Startet eine Lernepisode die dann endet, wenn entweder die angegeben Anzahl Schritte ausgeführt wurde oder eines
        /// der spezifizierten Zustände erreicht wurde. Es wird jedoch mindestens ein Schritt durchgeführt.<br />
        /// </summary>
        /// <param name="maxActionCount">Anzahl maximal auszuführender Schritte</param>
        /// <param name="finalStates">Endzustände</param>
        /// <exception cref="System.NullReferenceException" />
        /// <returns>Anzahl ausgeführter Schritte</returns>
        public override int RunEpisode(int maxActionCount, params S[] finalStates) {
            // alles Initialisiert?
            A action = Policy.GetAction( QFunction.GetEstimates(Environment.State) );
            int stepCount = 0;
            while ((!finalStates.Contains<S>(Environment.State)) && stepCount < maxActionCount) {
                action = this.DoStep(action);
                stepCount++;
            }
            // Falls einer der Endzustände erreicht wurde, wird die Lernepisode automatisch beendet.
            //if (finals.Contains<int>(environment.State))
            //    resetEpsiode();
            return stepCount;
        }

        /// <summary>
        /// Resets this instance.<para>
        /// Setzt EFunction und RewardTotal zurück, sodass eine neue Episode gestartet werden kann. 
        /// Das Erlernte (die QFunction) bleibt erhalten.</para>
        /// </summary>
        public override void Reset() {
            EFunction.Clear();
            RewardTotal = 0.0;
        }

        #endregion








        /// <summary>
        /// Standard Implementierung der E-Funktion. Wrapper für eine Dictionary. Initialer / Standard Wert für alle (s,a) ist 0.
        /// </summary>
        /// <remarks>
        /// Die <c>DefaultEFunc</c> ist für diskrete Zustandsräume. Für bspw. Linear Gradient descent ist die Funkntion ungeeignet!
        /// <para>
        /// <c>Path</c> wird genutzt, um die sa-Tupel zu speichern, der e-Wert nicht 0 ist. Es ist die "Spur" der letzten Zustände
        /// Aktionen, die zur aktuellen Belohnung beigetragen haben und somit für das Update zu berücksichtigen sind. Die QFunction
        /// kann Path ebenfalls für ein performantes Update nutzen.
        /// </para>
        /// </remarks>
        public class DefaultEFunc : IEligibilityFunction<A, S>
        {
            #region private Felder

            /// <summary>Hält die e-Wert (value) zu einem SA-Tupel (key).</summary>
            Dictionary<SATupel<S, A>, double> _eWerte;

            double _updFac;

            #endregion

            /// <summary>
            /// Gets or sets the path.
            /// Used for performance optimisation.
            /// </summary>
            /// <value>
            /// The path.
            /// </value>
            public TraceSet<SATupel<S, A>> Path { get; protected set; }


            #region Konstruktor(en)

            /// <summary>
            /// Initializes a new instance of the <see cref="EFuncAMEE"/> class.<para>
            /// Standardwerte: MinE = 0, TraceDecay = DiscountFactor = 1.</para>
            /// </summary>
            public DefaultEFunc() {
                Path = new TraceSet<SATupel<S, A>>();
                _eWerte = new Dictionary<SATupel<S, A>, double>();
                _gamma = 1.0;
                _lambda = 1.0;
            }

            #endregion


            #region IEligibilityFunction<A,S> Members

            /// <summary>
            /// Gets or sets the <see cref="System.Double"/> eligibility value for the specified state, action.
            /// </summary>
            /// <value>
            /// Der Wert des Berechtungsfaktors e für das spezifizierte sa-Tupel. Sollte in [0.0, 1.0] liegen.
            /// Wenn Wert = 0, dann wird sa-Tupel aus <c>Path</c> entfernt.
            /// </value>
            /// <param name="state">Zustand</param>
            /// <param name="action">Aktion</param>
            /// <returns>Den Berechtigungsfaktor e für das angegebene Tupel (state,action). Standardwert ist 0.0</returns>
            public double this[A action, S state] {
                get {
                    double retVal = 0.0;
                    try {
                        retVal = _eWerte[new SATupel<S, A>(state, action)];
                    }
                    catch (Exception) { /*nix*/
                        return retVal;
                    }
                    return retVal;
                }
                set {
                    var key = new SATupel<S, A>(state, action);
                    if (value == 0.0) {
                        _eWerte.Remove(key);
                        //Path.Remove(key);
                    }
                    else {
                        //Path.Add(key);
                        _eWerte[key] = value;
                    }
                }
            }


            /// <summary>
            /// Aktualisiert die E-Funktion mit ersetzenden Spuren. Fügt dem <c>Path</c> das
            /// angebene Aktions-Zustands-Tupel hinzu. 
            /// </summary>
            /// <param name="lastAction">The action.</param>
            /// <param name="lastState">The state.</param>
            /// <returns>This updated E-Function.</returns>
            public IEligibilityFunction<A, S> Update(A lastAction, S lastState) {
                //List<SATupel<S, A>> keys = new List<SATupel<S, A>>(_eWerte.Keys);
                Path.Add(new SATupel<S, A>(lastState, lastAction));
                var trace = new List<SATupel<S, A>>(Path); // Kopie, damit Path geändert werden kann
                foreach (var sa in trace) {
                    if (sa.State.Equals(lastState)) {
                        if (sa.Action.Equals(lastAction)) {
                            _eWerte[sa] = 1.0;
                        }
                        else {
                            //_eWerte[sa] = 0.0;
                            _eWerte.Remove(sa);
                            Path.Remove(sa); // !! Entfern sa aus Path
                        }
                    }
                    else {
                        double newVal = _eWerte[sa] * _updFac;
                        _eWerte[sa] = newVal < _minE ? 0.0 : newVal;
                    }
                }

                return this;
            }

            #endregion


            #region IEligibilityFunction Members

            /// <summary>
            /// Löscht alle Einträge, bzw. setzt sie auf einen Standardwert zurück.
            /// </summary>
            public void Clear() {
                _eWerte.Clear();
                Path.Clear();
            }

            /// <summary>
            /// Gets or sets the <see cref="System.Double"/> eligibility value with the specified state, action.<para>
            /// Implementierung von <see cref="Bettzueche.RLLibrary.Functions.IEligibilityFunction"/>.</para>
            /// </summary>
            /// <value>
            /// Der Wert des Berechtungsfaktors e. Sollte in [0.0, 1.0] liegen.
            /// </value>
            /// <param name="state">Zustand</param>
            /// <param name="action">Aktion</param>
            /// <returns>Den Berechtigungsfaktor e für das angegebene Tupel (state,action). Standardwert ist 0.0</returns>
            /// <exception cref="System.NotSupportedException">
            /// Mindest ein Agrument-Typ wird nicht unterstützt: state is WorldImage, action is AMEEAction ?
            /// </exception>
            public double this[IAction action, IState state] {
                get {
                    if (state is S) {
                        if (action is A) {
                            var s = (S)state;
                            var a = (A)action;
                            return this[a, s];
                        }
                        
                    }
                    throw new NotSupportedException("Mindest ein Agrument-Typ wird nicht unterstützt.");
                }
                set {
                    if (state is S) {
                        if (action is A) {
                            var s = (S)state;
                            var a = (A)action;
                            this[a, s] = value;
                        }
                    }
                    throw new NotSupportedException("Mindest ein Agrument-Typ wird nicht unterstützt.");
                }
            }

            ///// <summary>
            ///// Gets or sets the eligibility values.
            ///// </summary>
            ///// <value>
            ///// The eligibilities.
            ///// </value>
            //public Matrix Eligibilities { 
            //    get; 
            //    set; 
            //}

            /// <summary>
            /// Gets or sets the discount factor ('gamma').
            /// (used by Update method)
            /// </summary>
            /// <value>
            /// The discount factor.
            /// </value>
            public double DiscountFactor {
                get { return _gamma; }
                set { 
                    _gamma =value;
                    _updFac = _lambda * _gamma;
                }
            }
            double _gamma;

            /// <summary>
            /// Gets or sets the trace decay ('lambda').
            /// <para>(Used by Update method)</para>
            /// </summary>
            /// <value>
            /// The trace decay.
            /// </value>
            public double TraceDecay {
                get { return _lambda; }
                set { 
                    _lambda = value;
                    _updFac = _lambda * _gamma;
                }
            }
            double _lambda;

            double _minE;
            /// <summary>
            /// <c>MinE</c> gibt den kleinsten Wert für
            /// 'eligibility' an, der nicht zu 0.0 interpretiert wird. Standard = 0.0.
            /// </summary>
            /// <value>
            /// Der minimale Verantwortlichkeitsfaktor e. Negative Werte werden auf 0.0 gesetzt.
            /// </value>
            public double MinE {
                get { return _minE; }
                set { _minE = Math.Max(0.0, value); }
            }

            public IEligibilityFunction Update(IAction action, IState state) {
                if (action is A) {
                    if (state is S) {
                        return this.Update((A)action, (S)state);
                    }
                }
                throw new NotSupportedException("Mindest ein Agrument-Typ wird nicht unterstützt.");
            }
            #endregion


            #region Operators

            /// <summary>
            /// Implements the operator *.
            /// </summary>
            /// <param name="skalar">The skalar.</param>
            /// <param name="eFunc">The e function.</param>
            /// <returns>
            /// The result of the operator.
            /// </returns>
            /// <exception cref="System.NullReferenceException">eFunc</exception>
            public static DefaultEFunc operator*(double skalar, DefaultEFunc eFunc) {
                DefaultEFunc result = new DefaultEFunc { DiscountFactor = eFunc.DiscountFactor, TraceDecay= eFunc.TraceDecay };
                result._eWerte = new Dictionary<SATupel<S, A>, double>();
                foreach (var key in eFunc._eWerte.Keys) {
                    result._eWerte[/*new SATupel<S, A>(key.Item1, key.Item2)*/key] = skalar * eFunc._eWerte[key];
                }
                return result;
            }

            /// <summary>
            /// Implements the operator *.
            /// </summary>
            /// <param name="eFunc">The e function.</param>
            /// <param name="skalar">The skalar.</param>
            /// <returns>
            /// The result of the operator.
            /// </returns>
            /// <exception cref="System.NullReferenceException">eFunc</exception>
            public static DefaultEFunc operator*(DefaultEFunc eFunc, double skalar) {
                return skalar * eFunc;
            }
            
            #endregion
        }

    }
}
